#define SOLAR_USE_CUDA_INTER_OP_DATA
#include <algorithm/analytical_model/algorithm/AllModels.h>
#include <algorithm/analytical_model/algorithm/cpu/GeometryUtils.h>
#include <algorithm/analytical_model/algorithm/cpu/VectorConvert.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/Data.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/heliostat/RectangleHeliostat.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/receiver/CylinderReceiver.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/receiver/RectangleReceiver.h>
#include <algorithm/analytical_model/algorithm/cuda/inter_op/utils/convert/AllConverts.h>
#include <algorithm/analytical_model/algorithm/dnn/params/AllParams.h>
#include <algorithm/analytical_model/algorithm/pipeline/AnalyticalPipeline.h>
#include <algorithm/analytical_model/algorithm/pipeline/AnalyticalPipelineParam.h>
#include <algorithm/analytical_model/algorithm/pipeline/DumpArgs.h>
#include <algorithm/analytical_model/convolution/AllParas.h>
#include <algorithm/analytical_model/io/json/converter/AllConverters.h>
#include <algorithm/analytical_model/io/json/transform/TransformJSON.h>
#include <algorithm/analytical_model/simplified/AllParas.h>
#include <algorithm/shared/utils/GenericVectorMath.h>
#include <utils/Utils.h>


// TODO: 分离参数计算与IO部分

namespace solar
{
    using SimplifiedModelPara = std::variant<HFLCALPara, IHFLCALPara, BHFLCALPara, NEGPara>;
    using ConvolutionModelPara = std::variant<UNIZARPara, HuangPara, ICauchyPara, UNIZARPara2>;
    using DnnBasedModelPara = std::variant<NEGPara, UNIZARPara2>;
    using SimplifiedModelParaTypeList =
        boost::mp11::mp_list<HFLCALPara, IHFLCALPara, BHFLCALPara, NEGPara>;
    using ConvolutionModelParaTypeList =
        boost::mp11::mp_list<UNIZARPara, HuangPara, ICauchyPara, UNIZARPara2>;
    using DnnBasedModelParaTypeList = boost::mp11::mp_list<NEGPara, UNIZARPara2>;
    using DnnBasedModelInputTypeList = boost::mp11::mp_list<dnn::NEGInput, dnn::UNIZAR2Input>;
    using DnnBasedModelOutputTypeList = boost::mp11::mp_list<dnn::NEGOutput, dnn::UNIZAR2Output>;
} // namespace solar

auto solar::dumpAnalyticalModelArgs(AnalyticalPipelineParam* params) -> boost::json::value
{
    // Assume there's only one receiver

    auto receiver_geometry = SOLAR_GET_VALUE(Receiver, geometry_, params->receivers_[0]);
    // auto receiver_center = solar::getCenter(receiver_geometry.get());
    // JSON中缺少相关信息（经纬度、时间等），暂时这么写
    // 计算代码见 SceneRendererAdapter::increaseTime()
    auto sun_position = SOLAR_GET_VALUE(Sun, sun_position_, params->sun_);
    const auto azimuth_degree = SOLAR_GET_VALUE(SunPosition, azimuth_, *sun_position);
    const auto elevation_degree = SOLAR_GET_VALUE(SunPosition, elevation_, *sun_position);
    Float3 sunray_dir = cpu::toFloat3(-angleToDirection(azimuth_degree, elevation_degree));
    std::vector<Float3> heliostat_image_plane_axis(params->heliostats_.size() * 3);
    std::vector<cuda::inter_op::RectangleHeliostat> rectangle_heliostats(
        params->heliostats_.size());
    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (int i = 0; i < params->heliostats_.size(); i++)
    {
        solar::cuda::inter_op::convertHeliostat(std::addressof(params->heliostats_[i]),
                                                std::addressof(rectangle_heliostats[i]));
        rectangle_heliostats[i].setRowAndColumn({1, 1});
        rectangle_heliostats[i].setBelongingGridId(0);
        auto heliostat_geometry = SOLAR_GET_VALUE(Heliostat, geometry_, params->heliostats_[i]);
        auto focus_center = getFocusCenter(heliostat_geometry.get(), receiver_geometry.get());
        rectangle_heliostats[i].setNormalAndRotate(cpu::toFloat3(focus_center), sunray_dir);
        std::tie(heliostat_image_plane_axis[i * 3], heliostat_image_plane_axis[i * 3 + 1],
                 heliostat_image_plane_axis[i * 3 + 2]) =
            rectangle_heliostats[i].getImagePlaneAxis();
    }

    std::vector<SimplifiedModelPara> simplified_model_paras;
    std::vector<ConvolutionModelPara> convolution_model_paras;
    if (params->model_type_ == ModelType::kSimplified)
    {
        simplified_model_paras.resize(params->heliostats_.size());
    }
    if (params->model_type_ == ModelType::kConvolution)
    {
        convolution_model_paras.resize(params->heliostats_.size());
    }
    bool use_dnn_based_model = false;
    static const auto compute_sigma = [&]<typename T>(const T*)
    {
        static_assert(!boost::mp11::mp_contains<DnnBasedModelPara, T>::value);
        SOLAR_OMP_PARALLEL_FOR_SIMD
        for (int i = 0; i < rectangle_heliostats.size(); i++)
        {
            convertHeliostat(std::addressof(rectangle_heliostats[i]),
                             std::addressof(params->heliostats_[i]));
            T current_model_para;
            solar::computeSigma(
                std::addressof(params->simulation_para_), std::addressof(params->scene_para_),
                std::addressof(params->heliostat_paras_[i]), std::addressof(params->common_para_),
                std::addressof(params->sun_), std::addressof(params->heliostats_[i]),
                std::addressof(params->receivers_[0]), std::addressof(current_model_para));
            if constexpr (boost::mp11::mp_contains<SimplifiedModelPara, T>::value)
            {
                simplified_model_paras[i] = current_model_para;
            }
            if constexpr (boost::mp11::mp_contains<ConvolutionModelPara, T>::value)
            {
                convolution_model_paras[i] = current_model_para;
            }
        }
    };

    static const auto compute_sigma_dnn = [&]<typename T>(const T*)
    {
        static_assert(boost::mp11::mp_contains<DnnBasedModelPara, T>::value);
        use_dnn_based_model = true;
        // constexpr auto type_index = boost::mp11::mp_find<DnnBasedModelParaTypeList, T>::value;
        // using model_output_type = boost::mp11::mp_at_c<DnnBasedModelOutputTypeList, type_index>;
        const auto compute_sigma_callback = [&](void*) {};
        FuncConverter<void(void*)>::func_ = compute_sigma_callback;
        ComputeSigmaCallback compute_sigma_callback_func = FuncConverter<void(void*)>::callback;
        SOLAR_OMP_PARALLEL_FOR_SIMD
        for (int i = 0; i < rectangle_heliostats.size(); i++)
        {
            convertHeliostat(std::addressof(rectangle_heliostats[i]),
                             std::addressof(params->heliostats_[i]));
        }
        std::vector<T> model_paras(params->heliostats_.size());
        solar::computeSigma(std::addressof(params->simulation_para_),
                            std::addressof(params->scene_para_), params->heliostat_paras_.data(),
                            std::addressof(params->common_para_), std::addressof(params->sun_),
                            params->heliostats_.data(), std::addressof(params->receivers_[0]),
                            params->heliostats_.size(), model_paras.data(),
                            compute_sigma_callback_func);
        SOLAR_OMP_PARALLEL_FOR_SIMD
        for (int i = 0; i < model_paras.size(); i++)
        {
            const auto current_model_para = model_paras[i];
            if constexpr (boost::mp11::mp_contains<SimplifiedModelPara, T>::value)
            {
                simplified_model_paras[i] = current_model_para;
            }
            if constexpr (boost::mp11::mp_contains<ConvolutionModelPara, T>::value)
            {
                convolution_model_paras[i] = current_model_para;
            }
        }
    };

    // demo only
    // JSON中暂时没有这些数据
    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (int i = 0; i < params->heliostat_paras_.size(); i++)
    {
        SOLAR_SET_VALUE(HeliostatPara, sigma_tracking_, Vec2(0.0, 0.0), params->heliostat_paras_[i])
        // SOLAR_SET_VALUE(HeliostatPara, sigma_slope_error_, Vec2(0.001, 0.001),
        //                 params->heliostat_paras_[i])
    }

    if (params->model_type_ == ModelType::kSimplified)
    {
        switch (std::get<SimplifiedModelType>(params->detailed_model_type_))
        {
        case SimplifiedModelType::kHFLCAL:
        {
            HFLCALPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case SimplifiedModelType::kiHFLCAL:
        {
            IHFLCALPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case SimplifiedModelType::kbHFLCAL:
        {
            BHFLCALPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case SimplifiedModelType::kNEG:
        {
            NEGPara* ptr = nullptr;
            compute_sigma_dnn(ptr);
            break;
        }
        }
    }
    if (params->model_type_ == ModelType::kConvolution)
    {
        switch (std::get<ConvolutionModelType>(params->detailed_model_type_))
        {
        case ConvolutionModelType::kUNIZAR:
        {
            UNIZARPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case ConvolutionModelType::kHuang:
        {
            HuangPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case ConvolutionModelType::kiCauchy:
        {
            ICauchyPara* ptr = nullptr;
            compute_sigma(ptr);
            break;
        }
        case ConvolutionModelType::kUNIZAR2:
        {
            UNIZARPara2* ptr = nullptr;
            compute_sigma_dnn(ptr);
            break;
        }
        }
    }

    static constexpr char sigma_t[] = "sigma_";
    static constexpr char sigma_x_t[] = "sigma_x_";
    static constexpr char sigma_y_t[] = "sigma_y_";
    static constexpr char x3_t[] = "x3_";
    static constexpr char x0_t[] = "x0_";
    static constexpr char y1_t[] = "y1_";
    static constexpr char y0_t[] = "y0_";

    static constexpr char diffusion_coeff_t[] = "diffusion_coeff_";
    static constexpr char flux_param_t[] = "flux_param_";
    static const auto get_single_sigma_params =
        [&]<typename T>(T* para) -> std::tuple<double, double, double>
    {
        const auto sigma = para->*solar::getParaPtr<T, sigma_t>();
        const auto flux_param = para->*solar::getParaPtr<T, flux_param_t>();
        return {sigma, sigma, flux_param};
    };
    static const auto get_double_sigma_params =
        [&]<typename T>(T* para) -> std::tuple<double, double, double>
    {
        const auto sigma_x = para->*solar::getParaPtr<T, sigma_x_t>();
        const auto sigma_y = para->*solar::getParaPtr<T, sigma_y_t>();
        const auto flux_param = para->*solar::getParaPtr<T, flux_param_t>();
        return {sigma_x, sigma_y, flux_param};
    };
    static const auto get_convolution_area_rectangle =
        [&]<typename T>(T* para) -> std::tuple<double, double, double, double>
    {
        const auto x3 = para->*solar::getParaPtr<T, x3_t>();
        const auto x0 = para->*solar::getParaPtr<T, x0_t>();
        const auto y1 = para->*solar::getParaPtr<T, y1_t>();
        const auto y0 = para->*solar::getParaPtr<T, y0_t>();
        // 右上 左上 左下 左上
        return {x3, x0, y1, y0};
    };

    // returns (sigma_x, sigma_y, flux_param)
    static const auto get_sigmas = [&](int index) -> std::tuple<double, double, double>
    {
        if (params->model_type_ == ModelType::kSimplified)
        {
            if (auto* hflcal_para =
                    std::get_if<HFLCALPara>(std::addressof(simplified_model_paras[index]));
                hflcal_para != nullptr)
            {
                return get_single_sigma_params(hflcal_para);
            }
            if (auto* ihflcal_para =
                    std::get_if<IHFLCALPara>(std::addressof(simplified_model_paras[index]));
                ihflcal_para != nullptr)
            {
                return get_double_sigma_params(ihflcal_para);
            }
            if (auto* bhflcal_para =
                    std::get_if<BHFLCALPara>(std::addressof(simplified_model_paras[index]));
                bhflcal_para != nullptr)
            {
                return get_double_sigma_params(bhflcal_para);
            }
            if (auto* neg_para =
                    std::get_if<NEGPara>(std::addressof(simplified_model_paras[index]));
                neg_para != nullptr)
            {
                return get_double_sigma_params(neg_para);
            }
        }

        if (params->model_type_ == ModelType::kConvolution)
        {
            if (auto* unizar_para =
                    std::get_if<UNIZARPara>(std::addressof(convolution_model_paras[index]));
                unizar_para != nullptr)
            {
                return get_single_sigma_params(unizar_para);
            }
            if (auto* huang_para =
                    std::get_if<HuangPara>(std::addressof(convolution_model_paras[index]));
                huang_para != nullptr)
            {
                return get_double_sigma_params(huang_para);
            }
            if (auto* icauchy_para =
                    std::get_if<ICauchyPara>(std::addressof(convolution_model_paras[index]));
                icauchy_para != nullptr)
            {
                const auto diffusion_coeff =
                    icauchy_para->*solar::getParaPtr<ICauchyPara, diffusion_coeff_t>();
                const auto flux_param =
                    icauchy_para->*solar::getParaPtr<ICauchyPara, flux_param_t>();
                return {diffusion_coeff, diffusion_coeff, flux_param};
            }
            if (auto* unizar2_para =
                    std::get_if<UNIZARPara2>(std::addressof(convolution_model_paras[index]));
                unizar2_para != nullptr)
            {
                return get_double_sigma_params(unizar2_para);
            }
        }
        return {};
    };

    static const auto get_convolution_area =
        [&](int index) -> std::tuple<double, double, double, double>
    {
        if (params->model_type_ == ModelType::kConvolution)
        {
            if (auto* unizar2_para =
                    std::get_if<UNIZARPara2>(std::addressof(convolution_model_paras[index]));
                unizar2_para != nullptr)
            {
                return get_convolution_area_rectangle(unizar2_para);
            }
        }
        return {};
    };

    std::vector<Float3> flux_params(rectangle_heliostats.size());

    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (int i = 0; i < rectangle_heliostats.size(); i++)
    {
        const auto [sigma1, sigma2, peak_flux_param] = get_sigmas(i);
        flux_params[i] = Float3(sigma1, sigma2, peak_flux_param);
    }

    boost::json::value root_value;
    auto& root_object = root_value.emplace_object();
    boost::json::array heliostat_array;
    heliostat_array.resize(params->heliostats_.size());
    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (std::size_t i = 0; i < params->heliostats_.size(); i++)
    {
        auto& heliostat_value = heliostat_array[i];
        auto& heliostat_object = heliostat_value.emplace_object();
        const auto current_flux_param = flux_params[i];
        heliostat_object.insert_or_assign("peak_flux_param", current_flux_param.z);
        boost::json::array sigma_array;
        sigma_array.emplace_back(current_flux_param.x);
        sigma_array.emplace_back(current_flux_param.y);
        heliostat_object.insert_or_assign("sigmas", sigma_array);
    }

    // UNIZAR specific part
    // 需要导出定日镜在成像平面上的投影坐标

    SOLAR_OMP_PARALLEL_FOR_SIMD
    for (std::size_t index = 0; index < params->heliostats_.size(); index++)
    {
        if (params->model_type_ == ModelType::kConvolution)
        {
            auto& heliostat_value = heliostat_array[index];
            auto& heliostat_object = heliostat_value.get_object();
            boost::json::value unizar_extra_value;
            auto& unizar_extra_object = unizar_extra_value.emplace_object();
            const auto* current_image_plane_axis_ptr =
                heliostat_image_plane_axis.data() + (index * 3);
            const auto current_image_plane_x_axis = current_image_plane_axis_ptr[0];
            const auto current_image_plane_y_axis = current_image_plane_axis_ptr[1];
            // const auto current_image_plane_z_axis = current_image_plane_axis_ptr[2];
            const auto current_vertices = rectangle_heliostats[index].getVertices();
            const auto current_focus_center = rectangle_heliostats[index].getFocusCenter();
            const auto current_heliostat_verts = rectangle_heliostats[index].getVertices();

            // Float3 interv[4];
            // for (int i = 0; i < 4; i++)
            // {
            //     std::tie(interv[i], std::ignore) =
            //         cpu::intersect(current_image_plane_z_axis, current_focus_center,
            //                        current_heliostat_verts[i],
            //                        solar::internal::vector::negate(current_image_plane_z_axis));
            // }
            Float2 projected_interv[4];
            const auto current_focus_center_vec = cpu::toVec3(current_focus_center);
            for (int i = 0; i < 4; i++)
            {
                const auto current_heliostat_verts_i = cpu::toVec3(current_heliostat_verts[i]);
                const auto delta = current_heliostat_verts_i - current_focus_center_vec;
                projected_interv[i] =
                    Float2(solar::dot(cpu::toVec3(current_image_plane_x_axis), delta),
                           solar::dot(cpu::toVec3(current_image_plane_y_axis), delta));
            }
            Float2 current_projected_helio_verts[4];
            current_projected_helio_verts[0] = projected_interv[1];
            current_projected_helio_verts[1] = projected_interv[0];
            current_projected_helio_verts[2] = projected_interv[3];
            current_projected_helio_verts[3] = projected_interv[2];

            float x13 = current_projected_helio_verts[2].x;
            float x10 = current_projected_helio_verts[1].x;
            float y11 = current_projected_helio_verts[1].y;
            float y10 = current_projected_helio_verts[0].y;

            const auto projected_delta_y = y10 - y11;
            const auto projected_delta_x =
                current_projected_helio_verts[1].x - current_projected_helio_verts[0].x;
            const auto gradient_inv = projected_delta_x / projected_delta_y;

            if (use_dnn_based_model)
            {
                std::tie(x13, x10, y11, y10) = get_convolution_area(index);
            }

            unizar_extra_object.insert_or_assign("x13", x13);
            unizar_extra_object.insert_or_assign("x10", x10);
            unizar_extra_object.insert_or_assign("y11", y11);
            unizar_extra_object.insert_or_assign("y10", y10);
            unizar_extra_object.insert_or_assign("gradient_inv", gradient_inv);
            heliostat_object.insert_or_assign("UNIZAR", unizar_extra_value);
        }
    }

    root_object.insert_or_assign("heliostat_args", heliostat_array);

    return root_object;
}
